import torch
from torch.utils.data import Dataset
import pandas as pd
from PIL import Image
from torchvision import transforms
from . import custom_transforms as tr
from utils.process_labels import encode_labels,verify_labels

# 定义自己的类
class RoadLineDataset(Dataset):

    # 初始化
    def __init__(self, csv_path,base_size=16,is_train=True):
        self.base_size = base_size
        self.is_train = is_train
        # 读入数据
        fh = open(csv_path, 'r')
        self.imgs = []
        for line in fh:
            line = line.rstrip()
            pic_paths = line.split(",")
            self.imgs.append((pic_paths[0], pic_paths[1]))
    # 返回df的长度
    def __len__(self):
        return len(self.imgs)

    # 获取第idx+1列的数据
    def __getitem__(self, index):
        fn, label = self.imgs[index]
        input_img = Image.open(fn).convert('RGB')
        label_img = Image.open(label)
        sample = {'image': input_img, 'label': label_img}
        if self.is_train:
            transformed_result = self.transform_tr(sample)
        else:
            transformed_result = self.transform_val(sample)
        transformed_result['label'] = torch.from_numpy(encode_labels(transformed_result['label']))
        # label =  transformed_result['label'].numpy()
        return transformed_result
    def transform_tr(self, sample):
        composed_transforms = transforms.Compose([
            tr.RandomHorizontalFlip(),
            # tr.RandomScaleCrop(base_size=self.base_size, crop_size=self.base_size),
            tr.RandomGaussianBlur(),
            tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            tr.ToTensor()])

        return composed_transforms(sample)

    def transform_val(self, sample):

        composed_transforms = transforms.Compose([
            tr.FixScaleCrop(crop_size=513),
            tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            tr.ToTensor()])

        return composed_transforms(sample)
if __name__ == "__main__":
    ds = RoadLineDataset('train.csv')
    a= ds.__getitem__(5)